import cv2
import os
import numpy as np
#  添加了自己的position检测
#  加载人脸检测器
face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')


def hamming_distance(str1, str2):
    if len(str1) != len(str2):
        return
    count = 0
    for i in range(len(str1)):
        if str1[i] != str2[i]:
            count += 1
    return count


def dhash(image):
    #  转为灰度图像
    image = cv2.resize(image, (8, 8), interpolation=cv2.INTER_CUBIC)
    #  缩放为更小的尺寸方便计算
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    s = 0
    hash_str = ''
    for i in range(8):
        for j in range(8):
            s = s + gray[i, j]
    avg = s / 64
    for i in range(8):
        for j in range(8):
            if gray[i, j] > avg:
                hash_str = hash_str + '1'
            else:
                hash_str = hash_str + '0'
    return hash_str


def pos_bias(first_hash_value,x, y):#first_hash_value用来获取已经存在文件夹中的位置信息
    pos_bias_list = list()
    for i in range(len(faces_pos)):
        if faces_pos[i, 0]  == first_hash_value:
            pos_bias_value = abs(faces_pos[i, 1] - x) + abs(faces_pos[i, 2] - y)
            pos_bias_list.append(pos_bias_value)

    #     print(pos_bias_value)
    # print(pos_bias_list,min(pos_bias_list))

    return pos_bias_value


#  创建保存相似人脸的文件夹
def create_folder(path):
    if not os.path.exists(path):
        os.makedirs(path)


#  加载视频文件
video_path = r'D:\projects\dataset\bili.mp4'
cap = cv2.VideoCapture(video_path)

#  用于保存当前帧的人脸信息  {哈希值:  [人脸图像,  目标文件夹]}
faces_dict = {}
frame_count = 0

# 用于保存当前帧的位置信息 { 位置字典啊： [x,y]}

faces_pos = np.empty([0, 5], dtype=int)

while cap.isOpened():
    # ret, frame = cap.read()
    success, frame = cap.read()
    i = 0
    timeF = 40  # 修改timeF为你需要的间隔，假如视频是30帧/s， timeF=15就是0.5秒截图一次，timeF=30就是一秒截取一次
    j = 0  # 修改j=0可以设置命名的起始点，可以不修改(貌似是无用参数）
    while success:
        i = i + 1
        if (i % timeF == 0):
            j = j + 1
            print('Frame counter:', i)

            #  转换为灰度图像进行人脸检测
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

            #  人脸检测
            faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(60, 60))

            for (x, y, w, h) in faces:
                #  提取人脸图像
                face_img = frame[y:y + h, x:x + w]

                #  计算图像的哈希值
                hash_value = dhash(face_img)
                # print(x,y,w,h)

                if len(faces_dict) == 0:
                    #  创建第一个相似人脸文件夹
                    folder_path = os.path.join('D:/1/', str(hash_value))  # 设置保存相似人脸的文件夹路径
                    create_folder(folder_path)
                    face_path = os.path.join(folder_path, 'face{}.jpg'.format(frame_count))
                    # cv2.imwrite(face_path, face_img)
                    print("第一次获取hash值和位置信息")
                    # 获取的一个人脸的hash值
                    faces_dict[hash_value] = folder_path
                    # 获取人脸的位置信息
                    faces_pos = np.append(faces_pos, [[hash_value,x, y, w, h]], axis=0)
                    print(len(faces_dict),len(faces_pos))
                else:
                    #  判断是否有相似人脸已经存在
                    folder_paths = list(faces_dict.values())
                    for folder_path in folder_paths:
                        #  计算当前人脸图像与已存在文件夹中的人脸图像的汉明距离
                        face_paths = os.listdir(folder_path)
                        if len(face_paths) > 0:
                            first_face_path = os.path.join(folder_path, face_paths[0])
                            first_face_img = cv2.imread(first_face_path)
                            first_hash_value = dhash(first_face_img)
                            # 计算汉明距离
                            distance = hamming_distance(hash_value, first_hash_value)
                            # 计算人脸的位置
                            pos_value = pos_bias(first_hash_value,x, y)
                            if distance <= 5 or pos_value <= 50:  # 汉明距离小于等于5，人脸位置小于50，认为是相似人脸
                                face_path = os.path.join(folder_path, 'face{}.jpg'.format(frame_count))
                                cv2.imwrite(face_path, face_img)

                                # if i == 1000:
                                #     print("一千帧的时候人脸位置信息", faces_pos, len(faces_pos), len(faces_dict),
                                #           frame_count, faces_dict.items())

                                break
                    else:
                        #  与所有已存在文件夹中的人脸图像都不相似，创建新的文件夹
                        folder_path = os.path.join('D:/1/', str(hash_value))
                        create_folder(folder_path)
                        face_path = os.path.join(folder_path, 'face{}.jpg'.format(frame_count))
                        cv2.imwrite(face_path, face_img)
                        # 获取当前人脸的hash值
                        faces_dict[hash_value] = folder_path
                        # 增加人脸的位置信息
                        faces_pos = np.append(faces_pos, [[hash_value,x, y, w, h]], axis=0)
                        # print("再一次一次获取hash值和位置信息")
                        # print(len(faces_dict), len(faces_pos))

                frame_count += 1  # 检测到的人脸数

        # success, frame = cap.read()

    if not success:
        break

cap.release()
cv2.destroyAllWindows()
